import { FuncParameterType } from 'tinyest';
import { getAttributesString } from '../../data/attributes.ts';
import { type AnyData, undecorate } from '../../data/dataTypes.ts';
import {
  type ResolvedSnippet,
  snip,
  type Snippet,
} from '../../data/snippet.ts';
import {
  isPtr,
  isWgslData,
  isWgslStruct,
  Void,
  type WgslStruct,
} from '../../data/wgslTypes.ts';
import { MissingLinksError } from '../../errors.ts';
import { getMetaData, getName, setName } from '../../shared/meta.ts';
import type { ResolutionCtx } from '../../types.ts';
import {
  applyExternals,
  type ExternalMap,
  replaceExternalsInWgsl,
} from '../resolve/externals.ts';
import { extractArgs } from './extractArgs.ts';
import type { Implementation } from './fnTypes.ts';

export interface FnCore {
  applyExternals(newExternals: ExternalMap): void;
  resolve(
    ctx: ResolutionCtx,
    argTypes: AnyData[],
    /**
     * The return type of the function. If undefined, the type should be inferred
     * from the implementation (relevant for shellless functions).
     */
    returnType: AnyData | undefined,
  ): ResolvedSnippet;
}

export function createFnCore(
  implementation: Implementation,
  fnAttribute = '',
): FnCore {
  /**
   * External application has to be deferred until resolution because
   * some externals can reference the owner function which has not been
   * initialized yet (like when accessing the Output struct of a vertex
   * entry fn).
   */
  const externalsToApply: ExternalMap[] = [];

  const core = {
    applyExternals(newExternals: ExternalMap): void {
      externalsToApply.push(newExternals);
    },

    resolve(
      ctx: ResolutionCtx,
      argTypes: AnyData[],
      returnType: AnyData | undefined,
    ): ResolvedSnippet {
      const externalMap: ExternalMap = {};

      for (const externals of externalsToApply) {
        applyExternals(externalMap, externals);
      }

      const id = ctx.getUniqueName(this);

      if (typeof implementation === 'string') {
        if (!returnType) {
          throw new Error(
            'Explicit return type is required for string implementation',
          );
        }

        const replacedImpl = replaceExternalsInWgsl(
          ctx,
          externalMap,
          implementation,
        );

        let header = '';
        let body = '';

        if (fnAttribute !== '') {
          const input = isWgslStruct(argTypes[0])
            ? `(in: ${ctx.resolve(argTypes[0]).value})`
            : '()';

          const attributes = isWgslData(returnType)
            ? getAttributesString(returnType)
            : '';
          const output = returnType !== Void
            ? isWgslStruct(returnType)
              ? `-> ${ctx.resolve(returnType).value}`
              : `-> ${attributes !== '' ? attributes : '@location(0)'} ${
                ctx.resolve(returnType).value
              }`
            : '';

          header = `${input} ${output} `;
          body = replacedImpl;
        } else {
          const providedArgs = extractArgs(replacedImpl);

          if (providedArgs.args.length !== argTypes.length) {
            throw new Error(
              `WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${argTypes.length} arguments.`,
            );
          }

          const input = providedArgs.args.map((argInfo, i) =>
            `${argInfo.identifier}: ${
              checkAndReturnType(
                ctx,
                `parameter ${argInfo.identifier}`,
                argInfo.type,
                argTypes[i],
              )
            }`
          ).join(', ');

          const output = returnType === Void ? '' : `-> ${
            checkAndReturnType(
              ctx,
              'return type',
              providedArgs.ret?.type,
              returnType,
            )
          }`;

          header = `(${input}) ${output}`;

          body = replacedImpl.slice(providedArgs.range.end);
        }

        ctx.addDeclaration(`${fnAttribute}fn ${id}${header}${body}`);
        return snip(id, returnType, /* origin */ 'runtime');
      }

      // get data generated by the plugin
      const pluginData = getMetaData(implementation);

      // Passing a record happens prior to version 0.9.0
      // TODO: Support for this can be removed down the line
      const pluginExternals = typeof pluginData?.externals === 'function'
        ? pluginData.externals()
        : pluginData?.externals;

      if (pluginExternals) {
        const missing = Object.fromEntries(
          Object.entries(pluginExternals).filter(
            ([name]) => !(name in externalMap),
          ),
        );

        applyExternals(externalMap, missing);
      }

      const ast = pluginData?.ast;
      if (!ast) {
        throw new Error(
          "Missing metadata for tgpu.fn function body (either missing 'use gpu' directive, or misconfigured `unplugin-typegpu`)",
        );
      }

      // verify all required externals are present
      const missingExternals = ast.externalNames.filter(
        (name) => !(name in externalMap),
      );
      if (missingExternals.length > 0) {
        throw new MissingLinksError(getName(this), missingExternals);
      }

      // If an entrypoint implementation has a second argument, it represents the output schema.
      // We look at the identifier chosen by the user and add it to externals.
      const maybeSecondArg = ast.params[1];
      if (
        maybeSecondArg && maybeSecondArg.type === 'i' && fnAttribute !== ''
      ) {
        applyExternals(
          externalMap,
          {
            // biome-ignore lint/style/noNonNullAssertion: entry functions cannot be shellless
            [maybeSecondArg.name]: undecorate(returnType!),
          },
        );
      }

      // generate wgsl string

      const args: Snippet[] = [];
      const argAliases: [string, Snippet][] = [];

      for (const [i, argType] of argTypes.entries()) {
        const astParam = ast.params[i];
        // We know if arguments are passed by reference or by value, because we
        // enforce that based on the whether the argument is a pointer or not.
        //
        // It still applies for shell-less functions, since we determine the type
        // of the argument based on the argument's referentiality.
        // In other words, if we pass a reference to a function, it's typed as a pointer,
        // otherwise it's typed as a value.
        const origin = isPtr(argType)
          ? argType.addressSpace === 'storage'
            ? argType.access === 'read' ? 'readonly' : 'mutable'
            : argType.addressSpace
          : 'argument';

        switch (astParam?.type) {
          case FuncParameterType.identifier: {
            const rawName = astParam.name;
            const snippet = snip(ctx.makeNameValid(rawName), argType, origin);
            args.push(snippet);
            if (snippet.value !== rawName) {
              argAliases.push([rawName, snippet]);
            }
            break;
          }
          case FuncParameterType.destructuredObject: {
            args.push(snip(`_arg_${i}`, argType, origin));
            argAliases.push(...astParam.props.map(({ name, alias }) => {
              // Undecorating, as the struct type can contain builtins
              const destrType = undecorate(
                (argTypes[i] as WgslStruct).propTypes[name],
              );

              return [
                alias,
                snip(`_arg_${i}.${name}`, destrType, 'argument'),
              ] as [string, Snippet];
            }));
            break;
          }
          case undefined:
            args.push(snip(`_arg_${i}`, argType, origin));
        }
      }

      const { head, body, returnType: actualReturnType } = ctx.fnToWgsl({
        functionType: fnAttribute.includes('@compute')
          ? 'compute'
          : fnAttribute.includes('@vertex')
          ? 'vertex'
          : fnAttribute.includes('@fragment')
          ? 'fragment'
          : 'normal',
        args,
        argAliases: Object.fromEntries(argAliases),
        returnType,
        body: ast.body,
        externalMap,
      });

      ctx.addDeclaration(
        `${fnAttribute}fn ${id}${ctx.resolve(head).value}${
          ctx.resolve(body).value
        }`,
      );

      return snip(id, actualReturnType, /* origin */ 'runtime');
    },
  };

  // The implementation could have been given a name by a bundler plugin,
  // so we try to transfer it to the core.
  const maybeName = getName(implementation);
  if (maybeName !== undefined) {
    setName(core, maybeName);
  }

  return core;
}

function checkAndReturnType(
  ctx: ResolutionCtx,
  name: string,
  wgslType: string | undefined,
  jsType: unknown,
) {
  const resolvedJsType = ctx.resolve(jsType).value.replace(/\s/g, '');

  if (!wgslType) {
    return resolvedJsType;
  }

  const resolvedWgslType = wgslType.replace(/\s/g, '');

  if (resolvedWgslType !== resolvedJsType) {
    throw new Error(
      `Type mismatch between TGPU shell and WGSL code string: ${name}, JS type "${resolvedJsType}", WGSL type "${resolvedWgslType}".`,
    );
  }

  return wgslType;
}
